#!/usr/bin/env python3
"""
读取 ROS2 点云数据，分批绘制 IMU 世界系轨迹
Usage:
    python3 plot_imu_trajectory.py [--bag xxx.db3] [--bunch 20]
"""
import argparse
import numpy as np
import matplotlib.pyplot as plt
from matplotlib import animation
from rosbag2_py import SequentialReader, StorageOptions, ConverterOptions, StorageFilter
from sensor_msgs.msg import Imu
from rclpy.serialization import deserialize_message
from tqdm import tqdm

# 四元数转旋转矩阵
def quat2R(qw, qx, qy, qz):
    R = np.array([
        [1 - 2 * (qy ** 2 + qz ** 2), 2 * (qx * qy - qw * qz), 2 * (qx * qz + qw * qy)],
        [2 * (qx * qy + qw * qz), 1 - 2 * (qx ** 2 + qz ** 2), 2 * (qy * qz - qw * qx)],
        [2 * (qx * qz - qw * qy), 2 * (qy * qz + qw * qx), 1 - 2 * (qx ** 2 + qy ** 2)]
    ], dtype=np.float32)
    return R

# 读取并分批处理 IMU 数据
def load_imu_batches(bag_path, bunch_size):
    reader = SequentialReader()
    reader.open(StorageOptions(uri=bag_path, storage_id='sqlite3'), ConverterOptions('', ''))
    reader.set_filter(StorageFilter(topics=['/unilidar/imu']))

    batches = []
    cache_imu = []
    total_msgs = 0
    # 预估消息数（粗略，实际可用更精确方法）
    while reader.has_next():
        topic, _, _ = reader.read_next()
        if topic == '/unilidar/imu':
            total_msgs += 1
    # 重新打开reader
    reader = SequentialReader()
    reader.open(StorageOptions(uri=bag_path, storage_id='sqlite3'), ConverterOptions('', ''))
    reader.set_filter(StorageFilter(topics=['/unilidar/imu']))

    # 轨迹积分变量
    latest_R = np.eye(3, dtype=np.float32)
    latest_v = np.zeros(3, dtype=np.float32)  # 速度
    latest_t = np.zeros(3, dtype=np.float32)  # 位置
    last_time = None

    with tqdm(total=total_msgs, desc='加载IMU数据') as pbar:
        while reader.has_next():
            topic, data, t = reader.read_next()
            if topic == '/unilidar/imu':
                imu = deserialize_message(data, Imu)
                q = imu.orientation
                latest_R = quat2R(q.w, q.x, q.y, q.z)
                # 加速度（body系）
                acc_body = np.array([imu.linear_acceleration.x,
                                     imu.linear_acceleration.y,
                                     imu.linear_acceleration.z], dtype=np.float32)
                # 转到世界系
                acc_world = latest_R @ acc_body
                # 时间戳（纳秒转秒）
                curr_time = imu.header.stamp.sec + imu.header.stamp.nanosec * 1e-9 if hasattr(imu, 'header') else t * 1e-9
                if last_time is None:
                    dt = 0.01  # 默认步长
                else:
                    dt = curr_time - last_time
                last_time = curr_time
                # 积分速度和位置
                latest_v += acc_world * dt
                latest_t += latest_v * dt
                world_pos = latest_t.copy()
                cache_imu.append(world_pos)
                pbar.update(1)
                if len(cache_imu) == bunch_size:
                    batches.append(np.vstack(cache_imu))
                    cache_imu = []
        if cache_imu:
            batches.append(np.vstack(cache_imu))
    return batches


# 动画显示轨迹
def animate_imu_batches(batches, interval=50):
    fig = plt.figure(figsize=(8, 6))
    ax = fig.add_subplot(111, projection='3d')
    ax.set_xlabel('X world')
    ax.set_ylabel('Y world')
    ax.set_zlabel('Z world')
    ax.set_title('IMU Trajectory (World Coordinates)')
    scat = ax.scatter([], [], [], s=10, c='red')

    # 计算全局范围
    all_pts = np.vstack(batches) if batches else np.empty((0, 3))
    if all_pts.shape[0] > 0:
        margin = 1.0
        mins, maxs = all_pts.min(axis=0) - margin, all_pts.max(axis=0) + margin
        ax.set_xlim(mins[0], maxs[0])
        ax.set_ylim(mins[1], maxs[1])
        ax.set_zlim(mins[2], maxs[2])

    coords = [np.vstack(batches[:i+1]) for i in range(len(batches))]

    # tqdm 进度条
    pbar = tqdm(total=len(batches), desc='动画进度')

    def update(i):
        pts = coords[i]
        scat._offsets3d = (pts[:, 0], pts[:, 1], pts[:, 2])
        ax.set_title(f'IMU Trajectory - Step {i+1}/{len(batches)}')
        pbar.n = i + 1
        pbar.refresh()
        if i + 1 == len(batches):
            pbar.close()
        return scat,

    ani = animation.FuncAnimation(fig, update, frames=len(batches), interval=interval, blit=False, repeat=False)
    plt.show()

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--bag', default='/gitlab/pcd_analysis/demo_data/unilidar-2023-09-22-12-42-04.db3', help='ros2 bag db3 path')
    parser.add_argument('--bunch', type=int, default=100, help='frames per batch')
    parser.add_argument('--interval', type=int, default=10, help='ms per step')
    args = parser.parse_args()

    print('读取 bag 并分批加载 IMU 世界系轨迹 …')
    batches = load_imu_batches(args.bag, args.bunch)
    print(f'共 {len(batches)} 批')
    animate_imu_batches(batches, interval=args.interval)
